import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_dataset
import pandas as pd

# 加载预训练模型和分词器
model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 加载数据集
def load_data(file_path):
    dataset = pd.read_csv(file_path)
    conversations = []
    for _, row in dataset.iterrows():
        conversation = f"User: {row['user']}\nBot: {row['bot']}\n"
        conversations.append(conversation)
    return conversations

# 将数据转换为模型输入格式
def tokenize_function(examples):
    return tokenizer(examples, padding="max_length", truncation=True, max_length=128)

# 加载和准备数据
file_path = "dialogue_data.csv"
conversations = load_data(file_path)
dataset = load_dataset("text", data_files={"train": file_path})
tokenized_datasets = dataset.map(tokenize_function, batched=True)

# 定义训练参数
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
)

# 定义Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
)

# 训练模型
trainer.train()

# 保存模型
model.save_pretrained("./my_finetuned_model")
tokenizer.save_pretrained("./my_finetuned_model")